import cv2
import numpy as np
import os
import math
import concurrent.futures
from functools import partial

# --- 全局变量用于鼠标回调 (从原脚本复制，但在此脚本中不使用) ---
line_points = []
calibration_image_display = None

# --- Matplotlib 中文和负号显示配置 (从原脚本复制，但在此脚本中不使用) ---
# import matplotlib.pyplot as plt
# plt.rcParams['font.sans-serif'] = ['SimHei']
# plt.rcParams['axes.unicode_minus'] = False

# 从原脚本复制的核心图像处理函数
def process_single_image(file_info, input_folder, output_folder):
    """处理单个图像文件的函数，设计为在线程池中运行。"""
    i, filename, should_save = file_info # should_save 在此脚本中总是 True
    file_path = os.path.join(input_folder, filename)
    image = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)

    if image is None:
        print(f"警告: 无法读取图像 {file_path}")
        return None

    if len(image.shape) == 2:
        gray_image = image
    else:
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
    blurred_image = cv2.GaussianBlur(gray_image, (5, 5), 0)
    # 使用Otsu's阈值法自动确定最佳阈值
    _, binary_image = cv2.threshold(blurred_image, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    
    # 形态学操作：开运算去除小噪声，闭运算连接断裂部分
    kernel = np.ones((5, 5), np.uint8)
    binary_image = cv2.morphologyEx(binary_image, cv2.MORPH_OPEN, kernel, iterations=2)
    binary_image = cv2.morphologyEx(binary_image, cv2.MORPH_CLOSE, kernel, iterations=3)
    
    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if contours:
        largest_contour = max(contours, key=cv2.contourArea)
        M = cv2.moments(largest_contour)
        
        if M["m00"] != 0:
            cX_float = M["m10"] / M["m00"]
            cY_float = M["m01"] / M["m00"]
        else:
            # 如果m00为0，说明轮廓面积为0，退化为边界框中心
            x, y, w, h = cv2.boundingRect(largest_contour)
            cX_float, cY_float = float(x + w // 2), float(y + h // 2)

        # 始终保存图片并标注
        if len(image.shape) == 2:
            image_for_drawing = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
        else:
            image_for_drawing = image.copy()
        
        draw_cX, draw_cY = int(round(cX_float)), int(round(cY_float))
        
        cv2.drawContours(image_for_drawing, [largest_contour], -1, (0, 255, 0), 2) # 绿色轮廓
        cv2.circle(image_for_drawing, (draw_cX, draw_cY), 7, (0, 0, 255), -1) # 红色质心
        cv2.putText(image_for_drawing, f"({draw_cX}, {draw_cY})", (draw_cX + 10, draw_cY + 10),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) # 蓝色文字
        
        output_path = os.path.join(output_folder, filename)
        cv2.imwrite(output_path, image_for_drawing)
        
        return (filename, cX_float, cY_float) # 返回文件名和质心坐标，尽管在此脚本中不写入CSV
        
    return None

# --- 主程序 ---
if __name__ == "__main__":
    root_input_dir = input("请输入包含所有实验工况的根文件夹地址: ")
    if not os.path.isdir(root_input_dir):
        print(f"错误: 文件夹 '{root_input_dir}' 不存在。")
        exit()

    image_extensions = {'.bmp', '.png', '.jpg', '.jpeg'}

    # 查找所有包含图像文件的 "叶子" 实验文件夹 (与原脚本逻辑相同)
    experiment_folders = []
    for dirpath, dirnames, filenames in os.walk(root_input_dir):
        if os.path.basename(dirpath) in ['biaoding', 'bangzhongxin']:
            continue

        if any(os.path.splitext(f)[1].lower() in image_extensions for f in filenames):
            subdirs_with_images = False
            for sub_dir in dirnames:
                if sub_dir in ['biaoding', 'bangzhongxin']:
                    continue
                sub_dir_path = os.path.join(dirpath, sub_dir)
                if os.path.isdir(sub_dir_path):
                    for _, _, sub_files in os.walk(sub_dir_path):
                        if any(os.path.splitext(f)[1].lower() in image_extensions for f in sub_files):
                            subdirs_with_images = True
                            break
                    if subdirs_with_images:
                        break
            if not subdirs_with_images:
                experiment_folders.append(dirpath)

    if not experiment_folders:
        print(f"在 '{root_input_dir}' 及其子文件夹中未找到任何有效的实验文件夹。")
        exit()
    
    print(f"\n找到了 {len(experiment_folders)} 个实验文件夹，准备生成样图...")
    for folder in experiment_folders:
        print(f"  - {os.path.relpath(folder, root_input_dir)}")

    # 创建总输出目录
    sample_output_dir = os.path.join(os.path.dirname(root_input_dir.rstrip('/\\')), os.path.basename(root_input_dir.rstrip('/\\')) + "_sampled_centroids")
    os.makedirs(sample_output_dir, exist_ok=True)
    print(f"\n所有样图将保存在: {sample_output_dir}")

    TOTAL_SAMPLES_TO_GENERATE = 3000 # 总共需要生成的样图数量

    for i, input_folder in enumerate(experiment_folders):
        relative_path = os.path.relpath(input_folder, root_input_dir)
        current_output_folder = os.path.join(sample_output_dir, relative_path)
        os.makedirs(current_output_folder, exist_ok=True)

        print(f"\n--- [{i+1}/{len(experiment_folders)}] 正在为实验 '{relative_path}' 生成样图 ---")
        
        file_list = sorted([f for f in os.listdir(input_folder) if f.lower().endswith(tuple(image_extensions))])
        
        if not file_list:
            print(f"  - 警告: 文件夹 {input_folder} 中没有找到图像文件。跳过。")
            continue

        # 计算每个文件夹需要抽样的图片数量，确保总数接近3000
        num_files_in_folder = len(file_list)
        # 简单平均分配，如果文件夹数量多，每个文件夹可能抽样很少
        # 更精确的做法是按比例分配，但为了简化，这里直接除以文件夹数量
        
        # 确保每个文件夹至少有一张图片被抽样 (如果文件足够的话)
        num_samples_for_current_folder = max(1, math.ceil(TOTAL_SAMPLES_TO_GENERATE / len(experiment_folders)))
        
        # 确保不会超出当前文件夹的图片总数
        num_samples_for_current_folder = min(num_samples_for_current_folder, num_files_in_folder)


        # 等距抽样
        if num_samples_for_current_folder > 0:
            sample_indices = set(np.linspace(0, num_files_in_folder - 1, num_samples_for_current_folder, dtype=int))
        else:
            sample_indices = set()

        tasks = [(j, filename, True) for j, filename in enumerate(file_list) if j in sample_indices]
        
        if not tasks:
            print(f"  - 未从文件夹 '{relative_path}' 中抽取任何图片进行标注。")
            continue

        worker_func = partial(process_single_image, input_folder=input_folder, output_folder=current_output_folder)
        
        processed_count = 0
        try:
            from tqdm import tqdm
            with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
                for res in tqdm(executor.map(worker_func, tasks), total=len(tasks), desc=f"标注 {os.path.basename(input_folder)} 样图"):
                    if res is not None:
                        processed_count += 1
        except ImportError:
            print(f"正在标注文件夹: {input_folder} 的 {len(tasks)} 张样图...")
            with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
                for res in executor.map(worker_func, tasks):
                    if res is not None:
                        processed_count += 1
            print("标注完成。")

        print(f"  - 为 '{relative_path}' 生成了 {processed_count} 张标注样图，保存在: {current_output_folder}")

    print("\n\n所有样图生成完毕！请检查输出文件夹以验证形心识别的准确性。")